package common

import (
	"bufio"
	"context"
	"crypto/md5"
	"fmt"
	"io"
	"os"
	"path/filepath"
	"sync"
	"time"

	"go.uber.org/atomic"
)

const (
	ChunkSize      = 1024 * 1024 // 1MB 每块大小
	maxConcurrency = 5           // 最大并发处理数
	maxRetries     = 3           // 最大重试次数
	retryDelay     = time.Second // 重试延迟
)

type FileTransfer struct {
	mu         sync.Mutex
	files      map[string]*FileInfo
	tempDir    string
	workerPool chan struct{}   // 工作池限制并发
	ctx        context.Context // 用于控制和取消操作
	cancel     context.CancelFunc
}

type FileInfo struct {
	ID          string
	Size        int64
	Chunks      sync.Map // 使用sync.Map存储块数据
	ChunkStatus sync.Map // 记录块状态
	TotalChunks int
	Completed   *atomic.Bool
	Path        atomic.Value
	CreatedAt   time.Time
}

func NewFileTransfer(tempDir string) (*FileTransfer, error) {
	if err := os.MkdirAll(tempDir, 0755); err != nil {
		return nil, err
	}

	ctx, cancel := context.WithCancel(context.Background())
	ft := &FileTransfer{
		files:      make(map[string]*FileInfo),
		tempDir:    tempDir,
		workerPool: make(chan struct{}, maxConcurrency),
		ctx:        ctx,
		cancel:     cancel,
	}

	go ft.cleanupRoutine()
	return ft, nil
}

func (ft *FileTransfer) AddChunk(chunk *FileChunk) error {
	select {
	case ft.workerPool <- struct{}{}: // 获取工作槽
		defer func() { <-ft.workerPool }() // 释放工作槽
	case <-ft.ctx.Done():
		return fmt.Errorf("file transfer is shutting down")
	}

	return ft.processChunk(chunk)
}

func (ft *FileTransfer) processChunk(chunk *FileChunk) error {
	ft.mu.Lock()
	info, exists := ft.files[chunk.FileID]
	if !exists {
		info = &FileInfo{
			ID:          chunk.FileID,
			TotalChunks: chunk.TotalChunks,
			CreatedAt:   time.Now(),
			Completed:   atomic.NewBool(false),
		}
		ft.files[chunk.FileID] = info
	}
	ft.mu.Unlock()

	// 使用原子操作记录块状态
	if !info.setChunkStatus(chunk.ChunkIndex) {
		return nil // 块已经处理过
	}

	// 保存块数据
	info.setChunk(chunk.ChunkIndex, chunk.Data)

	// 检查是否所有块都已接收
	if info.isComplete() {
		return ft.assembleFile(info)
	}

	return nil
}

func (ft *FileTransfer) assembleFile(info *FileInfo) error {
	// 创建临时文件
	tempFile := filepath.Join(ft.tempDir, info.ID)
	file, err := os.OpenFile(tempFile, os.O_CREATE|os.O_WRONLY, 0644)
	if err != nil {
		return fmt.Errorf("create file error: %v", err)
	}
	defer file.Close()

	// 使用bufio提升写入性能
	writer := bufio.NewWriter(file)

	// 按顺序写入块
	for i := 0; i < info.TotalChunks; i++ {
		data := info.getChunk(i)
		if data == nil {
			return fmt.Errorf("missing chunk %d", i)
		}

		if _, err := writer.Write(data); err != nil {
			return fmt.Errorf("write chunk error: %v", err)
		}
	}

	if err := writer.Flush(); err != nil {
		return fmt.Errorf("flush error: %v", err)
	}

	info.setCompleted(tempFile)
	return nil
}

func (fi *FileInfo) setChunk(index int, data []byte) {
	fi.Chunks.Store(index, data)
}

func (fi *FileInfo) getChunk(index int) []byte {
	if data, ok := fi.Chunks.Load(index); ok {
		return data.([]byte)
	}
	return nil
}

func (fi *FileInfo) setChunkStatus(index int) bool {
	_, loaded := fi.ChunkStatus.LoadOrStore(index, true)
	return !loaded // 返回true表示这是新的块
}

func (fi *FileInfo) isComplete() bool {
	count := 0
	fi.ChunkStatus.Range(func(_, _ interface{}) bool {
		count++
		return true
	})
	return count == fi.TotalChunks
}

func (fi *FileInfo) setCompleted(path string) {
	fi.Completed.Store(true)
	fi.Path.Store(path)
}

func (ft *FileTransfer) GetFile(fileID string) (string, bool) {
	ft.mu.Lock()
	defer ft.mu.Unlock()

	if info, exists := ft.files[fileID]; exists && info.Completed.Load() {
		return info.Path.Load().(string), true
	}
	return "", false
}

func (ft *FileTransfer) cleanupRoutine() {
	ticker := time.NewTicker(time.Hour)
	defer ticker.Stop()

	for range ticker.C {
		ft.mu.Lock()
		now := time.Now()
		for id, info := range ft.files {
			// 清理超过24小时的文件
			if now.Sub(info.CreatedAt) > 24*time.Hour {
				if path, ok := info.Path.Load().(string); ok {
					os.Remove(path)
				}
				delete(ft.files, id)
			}
		}
		ft.mu.Unlock()
	}
}

// 生成文件ID
func GenerateFileID(path string, size int64) string {
	h := md5.New()
	io.WriteString(h, path)
	io.WriteString(h, fmt.Sprintf("%d", size))
	io.WriteString(h, time.Now().String())
	return fmt.Sprintf("%x", h.Sum(nil))
}

func (ft *FileTransfer) Close() error {
	ft.cancel() // 取消所有正在进行的操作
	close(ft.workerPool)

	ft.mu.Lock()
	defer ft.mu.Unlock()

	// 清理所有临时文件
	for _, info := range ft.files {
		if path, ok := info.Path.Load().(string); ok {
			os.Remove(path)
		}
	}

	return nil
}
